#include "Optimization.h"
#include "InterferenceAnalysis.h"
#include "ScopeDictionary.h"

namespace Compiler
{
	namespace Intermediate
	{
		Operation opParallelCopy(L"ParallelCopy");
		static Operation * ParallelCopy = &opParallelCopy;


		class OutOfSSA_Transform : public IntraProcOptimizer
		{
		private:
			template <typename F>
			void ForEachPhi(ControlFlowGraph * program, F func)
			{
				auto traverse = program->GetPostOrderOnReverseCFG();
				for (int i = traverse.Count()-1; i>=0; i--)
				{
					auto node = traverse[i];
					for (auto instrNode = node->Code.FirstNode(); instrNode; instrNode = instrNode->GetNext())
					{
						auto & instr = instrNode->Value;
						if (instr.Func != Operation::Phi) break;
						func(instrNode, instr);
					}
				}
			}
			template <typename F>
			void ForEachParallelCopy(ControlFlowGraph * program, F func)
			{
				for (int i = 0; i<program->Nodes.Count(); i++)
				{
					auto node = program->Nodes[i].Ptr();
					for (auto instrNode = node->Code.begin(); instrNode != node->Code.end(); ++instrNode)
					{
						auto & instr = instrNode.Current->Value;
						if (instr.Func != ParallelCopy) continue;
						func(instrNode.Current, instr);
					}
				}
			}
			
			struct TransformContext
			{
				List<RefPtr<Variable>> SSA_Vars;
				List<int> SSA_VarMap; // ssa var -> coalecsed var
				List<int> SSA_ValueNumber;
				List<RefPtr<EnumerableIntSet>> PhiClasses; // in ssa vars
				InterferenceGraph InterGraph; // in coalecsed var
				TransformContext()
				{}

				TransformContext(TransformContext && ctx)
				{
					PhiClasses = _Move(ctx.PhiClasses);
					InterGraph = _Move(ctx.InterGraph);
					SSA_VarMap = _Move(ctx.SSA_VarMap);
					SSA_Vars = _Move(ctx.SSA_Vars);
					SSA_ValueNumber = _Move(ctx.SSA_ValueNumber);
				}
			};

			StringBuilder nameBuilder;

			Variable * MakePhiCopy(ControlFlowGraph * program, Variable * oldVar, InstructionNode * defNode)
			{
				auto targetVar = new Variable(*oldVar);
				nameBuilder.Clear();
				nameBuilder<<L"^phi_"<<program->Variables.Count()<<L'_'<<targetVar->Name;
				targetVar->Name = nameBuilder.ToString();
				targetVar->Id = program->Variables.Count();
				program->Variables.Add(targetVar);
				program->VarDefs.Add(defNode);
				return targetVar;
			}

			void InsertCopies(ControlFlowGraph * program)
			{
				ForEachPhi(program, [&](InstructionNode * instrNode, Instruction & instr)
				{
					for (int i = 0; i < instr.Operands.Count(); i++)
					{
						auto node = instr.CFG_Node->Entries[i];
						InstructionNode * instrCopy = node->Code.LastNode()->GetPrevious();
						if (instrCopy == 0 ||
							instrCopy->Value.Func != ParallelCopy)
						{
							instrCopy = node->Code.LastNode()->InsertBefore(Instruction(ParallelCopy));
						}
						auto targetVar = MakePhiCopy(program, instr.Operands[i].Var, instrCopy);
						
						instrCopy->Value.Operands.Add(targetVar);
						instrCopy->Value.Operands.Add(instr.Operands[i].Var);
						instr.Operands[i].Var = targetVar;
					}
					// for left operand:
					{
						InstructionNode * instrCopy = instr.CFG_Node->FirstInstruction();
						if (!instrCopy)
						{
							instrCopy = instr.CFG_Node->Code.AddLast(Instruction(ParallelCopy));
						}
						else if (instrCopy->Value.Func != ParallelCopy)
						{
							instrCopy = instrCopy->InsertBefore(Instruction(ParallelCopy));
						}
						auto newVar = MakePhiCopy(program, instr.LeftOperand.Var, instrCopy);
						instrCopy->Value.Operands.Add(instr.LeftOperand.Var);
						instrCopy->Value.Operands.Add(newVar);
						instr.LeftOperand.Var = newVar;
					}
				});
			}

			TransformContext CoalescePhi(ControlFlowGraph * program)
			{
				TransformContext context;
				for (auto & var : program->Variables)
				{
					auto phiClass = new EnumerableIntSet();
					phiClass->Set.SetMax(program->Variables.Count());
					phiClass->Add(var->Id);
					context.PhiClasses.Add(phiClass);
				}
				ForEachPhi(program, [&](InstructionNode * instrNode, Instruction & instr)
				{
					auto phiClass = context.PhiClasses[instr.LeftOperand.Var->Id];
					for (auto & op : instr.Operands)
					{
						phiClass->UnionWith(*context.PhiClasses[op.Var->Id]);
						context.PhiClasses[op.Var->Id] = phiClass;
					}
				});
				context.SSA_Vars = program->Variables;
				context.SSA_VarMap.SetSize(program->Variables.Count());
				return context;
			}

			template <typename F>
			void ForEachCopyPair(Instruction & instr, F func)
			{
				for (int i = 0; i<instr.Operands.Count()>>1; i++)
				{
					if (instr.Operands[i*2].Var != instr.Operands[i*2+1].Var)
						func(instr.Operands[i*2].Var, instr.Operands[i*2+1].Var);
				}
			}

			void SequentializeParallelCopy(ControlFlowGraph * program)
			{
				Dictionary<DataType, Variable *> tempVars;
				List<Variable *> ready, todo;
				Dictionary<Variable *, Variable *> loc, pred;
				ForEachParallelCopy(program, [&](InstructionNode * instrNode, Instruction & instr)
				{
					ready.Clear();
					todo.Clear();
					loc.Clear();
					pred.Clear();
					ForEachCopyPair(instr, [&](Variable * b, Variable * a) // in the form of b = a
					{
						loc[0] = 0;
						pred[0] = 0;
						loc[b] = 0;
						pred[a] = 0;
					});
					ForEachCopyPair(instr, [&](Variable * b, Variable * a) // in the form of b = a
					{
						loc[a] = a;
						pred[b] = a;
						todo.Add(a);
					});
					ForEachCopyPair(instr, [&](Variable * b, Variable * a) // in the form of b = a
					{
						if (loc[b] == 0)
							ready.Add(b);
					});
					auto EmitCopy = [&](Variable * dest, Variable * src)
					{
						instrNode->InsertAfter(Instruction(dest, 0, src));
					};
					while (todo.Count())
					{
						Variable * a, * b, * c;
						while (ready.Count())
						{
							b = ready.Last();
							ready.RemoveAt(ready.Count()-1);
							a = pred[b];
							c = loc[a];
							EmitCopy(b, c);
							loc[a] = b;
							if (a== c && pred[a] != 0)
								ready.Add(a);
						}
						b = todo.Last();
						todo.RemoveAt(todo.Count()-1);
						if (b == loc[pred[b]])
						{
							Variable * tmpVar;
							if (!tempVars.TryGetValue(b->Type, tmpVar))
							{
								tmpVar = new Variable(*b);
								tmpVar->Name = L"^tmpcpy_" + String((int)b->Type);
								tmpVar->Version = 0;
								tempVars[b->Type] = tmpVar;
							}
							EmitCopy(tmpVar, b);
							loc[b] = tmpVar;
							ready.Add(b);
						}
					}
					instrNode->Delete();
				});
			}
			void Rename(ControlFlowGraph * program, TransformContext & context)
			{
				List<RefPtr<Variable>> newVars;
				Dictionary<Variable*, Variable*> varMapping;
				Dictionary<Variable*, Variable*> oldVarMapping;
				List<RefPtr<EnumerableIntSet>> newPhiClasses;
				for (int i = 0; i<program->Variables.Count(); i++)
				{
					auto curVar = program->Variables[i].Ptr();
					if (i < program->ParameterCount)
					{
						newVars.Add(program->Variables[i]);
						varMapping[program->Variables[i].Ptr()] = program->Variables[i].Ptr();
						oldVarMapping[program->Variables[i].Ptr()] = program->Variables[i].Ptr();
					}
					else
					{
						auto phiClass = context.PhiClasses[i];
						if (phiClass)
						{
							int representId = Min(phiClass->Values);
							Variable * representVar = context.SSA_Vars[representId].Ptr();
							Variable * newVar;
							if (!varMapping.TryGetValue(representVar, newVar))
							{
								newVar = new Variable(*representVar);
								if (representVar->Version > 0)
									newVar->Name = L"^phi_" + String(representVar->Version) + L"_" + newVar->Name; 
								newVar->Version = 0;
								newVar->Id = newVars.Count();
								newVars.Add(newVar);
							}
							oldVarMapping[newVar] = curVar;
							varMapping[curVar] = newVar;
							varMapping[representVar] = newVar;
						}
						else
						{
							auto newVar = new Variable(*curVar);
							newVar->Version = 0;
							newVar->Id = newVars.Count();
							newVars.Add(newVar);
							varMapping[curVar] = newVar;
							oldVarMapping[newVar] = curVar;
						}
					}
				}
				
				for (auto & node : program->Nodes)
				{
					for (auto instrNode = node->Code.begin(); instrNode!=node->Code.end(); ++instrNode)
					{
						auto & instr = instrNode.Current->Value;
						if (instr.Func == Operation::Phi)
						{
							instrNode.Current->Delete();
							continue;
						}
						
						if (instr.LeftOperand.IsVariable())
						{
							instr.LeftOperand.Var = varMapping[instr.LeftOperand.Var];
						}
						for (auto & op : instr.Operands)
							if (op.IsVariable())
								op.Var = varMapping[op.Var];
						if (instr.Func == 0 && instr.Operands[0].IsVariable() && instr.LeftOperand.Var == instr.Operands[0].Var)
						{
							instrNode.Current->Delete();
							continue;
						}
					}
				}
				for (auto var : newVars)
				{
					auto oldVar = oldVarMapping[var.Ptr()].GetValue();
					auto oldLiveRange = context.PhiClasses[oldVar->Id];
					newPhiClasses.Add(oldLiveRange);
					for (auto & ssaVar : oldLiveRange->Values)
						context.SSA_VarMap[ssaVar] = var->Id;
				}
				
				program->Variables = _Move(newVars);
				context.PhiClasses = _Move(newPhiClasses);
			}
			bool PhiClassInterferes(TransformContext & context, RefPtr<EnumerableIntSet> & phi1, RefPtr<EnumerableIntSet> & phi2)
			{
				for (int i = 0; i<phi1->Values.Count(); i++)
				{
					int var1 = context.SSA_VarMap[phi1->Values[i]];
					for (int j = 0; j<phi2->Values.Count(); j++)
					{
						int var2 = context.SSA_VarMap[phi2->Values[j]];
						if (var1 != var2 && context.InterGraph.Interferes(var1, var2)
							&& context.SSA_ValueNumber[var1] != context.SSA_ValueNumber[var2])
							return true;
					}
				}
				return false;
			}
			void TryCoalesceVariable(TransformContext & context, Variable * var1, Variable * var2)
			{
				auto & phi1 = context.PhiClasses[var1->Id];
				auto & phi2 = context.PhiClasses[var2->Id];
				if (phi1 != phi2 && !PhiClassInterferes(context, phi1, phi2))
				{
					phi1->UnionWith(*phi2);
					for (auto var : phi1->Values)
					{
						context.PhiClasses[context.SSA_VarMap[var]] = phi1;
					}
					for (auto var : phi2->Values)
					{
						context.PhiClasses[context.SSA_VarMap[var]] = phi1;
					}
				}
			}
			void CoalesceCopy(ControlFlowGraph * program, TransformContext & context)
			{
				for (auto & node : program->Nodes)
				{
					for (auto & instr : node->Code)
					{
						if (instr.Func == 0)
						{
							if (instr.Operands[0].IsVariable())
								TryCoalesceVariable(context, instr.LeftOperand.Var, instr.Operands[0].Var);
						}
						else if (instr.Func == ParallelCopy)
						{
							for (int i = 0; i<instr.Operands.Count()>>1; i++)
							{
								TryCoalesceVariable(context, instr.Operands[i*2].Var, instr.Operands[i*2+1].Var);
							}
						}
					}
				}
			}
			void EvalValueNumber(ControlFlowGraph * program, TransformContext & context)
			{
				context.SSA_ValueNumber.SetSize(context.SSA_Vars.Count());
				for (int i = 0; i<context.SSA_ValueNumber.Count(); i++)
					context.SSA_ValueNumber[i] = i;
				auto traverse = program->GetPostOrder();
				for (int i = traverse.Count()-1; i>=0; i--)
				{
					auto & node = traverse[i];
					for (auto & instr : node->Code)
					{
						if (instr.Func == 0 && instr.Operands[0].IsVariable())
							context.SSA_ValueNumber[instr.LeftOperand.Var->Id] = context.SSA_ValueNumber[instr.Operands[0].Var->Id];
						else if (instr.Func == ParallelCopy)
						{
							ForEachCopyPair(instr, [&](Variable *a, Variable * b)
							{
								context.SSA_ValueNumber[a->Id] = context.SSA_ValueNumber[b->Id];
							});
						}
					}
				}

			}
			void OutOfSSA(ControlFlowGraph * program)
			{
				InsertCopies(program);
				TransformContext context = CoalescePhi(program);
				EvalValueNumber(program, context);
				Rename(program, context);
				SequentializeParallelCopy(program);
				context.InterGraph = InterferenceAnalysis::BuildInterferenceGraph(program);
				/*for (auto node : program->Nodes)
				{
					printf("live out(%d): ", node->Id);
					for (int i = 0; i<program->Variables.Count(); i++)
					{
						if (node->LiveOut.Contains(i))
							printf("%d ", i);
					}
					printf("\n");
				}*/
				CoalesceCopy(program, context);
				Rename(program, context);
			}
		public:
			virtual ProgramOptimizationResult Optimize(RefPtr<ControlFlowGraph> program) override
			{
				ProgramOptimizationResult rs;
				rs.Changed = true;
				rs.Program = program;
				OutOfSSA(program.Ptr());
				return rs;
			}
		};

		IntraProcOptimizer * CreateOutOfSSA_Transform()
		{
			return new OutOfSSA_Transform();
		}
	}
}